import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import os
from tqdm import tqdm
from random_case import get_10_random_cases
from random_case import find_tsv_files
from random_case import find_cases
from collections import defaultdict
from datetime import datetime
import random

def draw_data_matrix(file_path):
    project = file_path.split('.')[0]
    # 声明：csv文件第一列为索引，第一行为列名
    fig = plt.figure()
    df=pd.read_csv(file_path,sep='\t',index_col=0,header=0) # 第一行第一列不读取
    # df2 = pd.DataFrame(df.values.T, index=df.columns, columns=df.index) # 转置，shape变成（594，31598）
    df2 = df.transpose()
    data = df2.values # DataFrame转numpy数组
    ax1 = sns.heatmap(data=data, xticklabels =5000) # 画热力图
    ax1.set_title(project)
    ax1.set_xlabel('RNA attribute')
    ax1.set_ylabel('samples')
    # plt.show()
    # 保存图片
    fig = ax1.get_figure()
    
    filename = f'数据矩阵图-2  -- {project}'
    output_folder = '1-data-matrix'
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    fig.savefig(os.path.join(output_folder, filename))

def pre_process(file_dir):
    '''
    函数简介：根据癌症工程所在的路径进行预处理
    '''
    temp_df = pd.DataFrame()
    cases = find_cases(file_dir)
    for i, case in enumerate(tqdm(cases)):
        df_test = pd.read_csv(case[0],sep='\t',comment='#',header=0,index_col=0)
        temp_df = pd.concat([temp_df,df_test.iloc[4:,7]],axis=1,ignore_index=False)
    # temp_df = temp_df.transpose()
    print(temp_df,temp_df.shape)
    # return
    index_list = []
    for i in range(temp_df.shape[0]):
        number = dict(temp_df.iloc[i].value_counts(normalize=True)) # 统计第i行所有元素的个数(以百分比形式显示)
        number_default_value = defaultdict(int,number) # 为防止字典里没有KEY0.0，从而使用defaultdict 
        if number_default_value[0.0] > 0.5: # 如果某RNA在一半以上的病例中没有表达（数值为0），则将其丢弃
            index_list.append(i) # 将对应的行号放入列表中
    # 删除指定行（index_list里的行数）
    temp_df = temp_df.drop(temp_df.index[index_list])
    temp_df = temp_df.apply(lambda x : np.log10(x+1))       #对数标准化
    new_column_names = [str(i) for i in range(temp_df.shape[1])]
    temp_df.columns = new_column_names
    print(temp_df,temp_df.shape)
    temp_df.to_csv(f'{file_dir}.tsv',sep='\t')
    return

def draw_hist_box(file_path):
    project = file_path.split('.')[0]
    df=pd.read_csv(file_path,sep='\t',index_col=0,header=0).transpose() # 第一行第一列不读取
    fig = plt.figure(figsize = (10,20))
    for i in range(10):
        # 读取所有样本的第三列属性
        attribute1 = df.iloc[:,i+1]
        # 画直方图
        fig.add_subplot(5,4,2*i+1)
        sns.histplot(attribute1)
        # 画箱线图
        fig.add_subplot(5,4,2*i+2)
        sns.boxplot(data=attribute1)
        # 设置总标题    
        plt.suptitle(f"{project}")
    plt.tight_layout(pad = 2.2)
    # plt.subplots_adjust(bottom=0.15)
    # plt.show()
    filename = f'直方图和盒状图 -- {project}'
    output_folder = '2-hist-box'
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    plt.savefig(os.path.join(output_folder, filename))

def draw_histogram(file_path = ""):
    rd_cases = get_10_random_cases(file_path)
    plt.figure(figsize=(16, 9))
    for i, case in enumerate(rd_cases):
        df_test = pd.read_csv(case,sep='\t',comment='#',index_col=None, usecols=['fpkm_uq_unstranded'],skiprows=[2,3,4,5])
        # df_test = df.iloc[4:,7]     #只取最后一列
        df_test = df_test.apply(lambda x : np.log10(x+1))       #对数标准化
        plt.subplot(5, 2, i+1)      #做子图
        sns.histplot(df_test, bins = 20, kde=False)
        plt.title(f'Histogram of Case: {case[1]}')
    plt.tight_layout(pad = 0.2)  #调整间距
    plt.show()
    pass



def draw_box(file_path = ""):
    rd_cases = get_10_random_cases(file_path)
    plt.figure(figsize=(16, 9))
    for i, case in enumerate(rd_cases):
        df_test = pd.read_csv(case,sep='\t',comment='#',index_col=None, usecols=['fpkm_uq_unstranded'],skiprows=[2,3,4,5])
        # df_test = df.iloc[4:,7]     #只取最后一列
        df_test = df_test.apply(lambda x : np.log10(x+1))       #对数标准化
        plt.subplot(5, 2, i+1)      #做子图
        sns.boxplot(df_test)
        plt.title(f'Box of Case: {case[1]}')
    plt.tight_layout(pad = 0.2)  #调整间距
    plt.show()

def draw_tsne_2d_map(file_path = ""):
    cases = find_cases(file_path)
    Project = cases[0][0].split('\\')[0]
    cases_sample_type = pd.DataFrame(cases).iloc[0:,1].values
    df_test = pd.read_csv(file_path+'.tsv',sep='\t',index_col=0,header=0)
    temp_df = df_test.transpose()
    for p in [5,25,50]:
        tsne = TSNE(n_components=2, random_state=0,perplexity=p,init='pca')
        reduced_data = tsne.fit_transform(temp_df)
        plt.figure(figsize=(10, 8))
        # 选取sample type作为标签来映射颜色
        sns.scatterplot(x=reduced_data[:, 0], y=reduced_data[:, 1], hue=cases_sample_type, palette='viridis')
        plt.title(f't-SNE 2D Visualization of : {Project} -- perplexity: {p}')
        # 构建文件名，时间格式 '2024-04-13-15-30-45.png'
        current_time = datetime.now().strftime('%Y-%m-%d-%H-%M-%S')
        filename = f't-SNE 2D Visualization of {Project}--perplexity-{p}-{current_time}.png'
        output_folder = f'tsne2d\\{Project}'
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        plt.savefig(os.path.join(output_folder, filename))
        plt.show()

if __name__ == '__main__':
    file_list=[
        'TCGA-BLCA',
        'TCGA-BRCA',
        'TCGA-LGG',
        'TCGA-LUAD',
        'TCGA-LUSC',
    ]
    for file_dir in file_list:
        draw_hist_box(file_path=file_dir+'.tsv')


    pass